/**
 * @file   Stokes.cpp
 * @author HirasawaYui <yui@Ubuntu18-04>
 * @date   Tue Jan 12 18:36:49 2021
 * 
 * @brief  一个有限元方法求解 Stokes 方程的算例
 * 
 * 
 */

//g++ -o main Stokes.cpp -std=c++11 -I /usr/include/eigen3/ ./include/tinyxml2.cpp -g

#include "Mesh.h"
#include "Element.h"
#include "MultigridSolver.h"
#include <Eigen/Sparse>
#include <cmath>
#include "Surface.h"
#include "MINRES.h"

#define pi 4 * atan(1.0)

typedef Eigen::SparseMatrix<Real> SpMat;
typedef Eigen::Triplet<Real> Tri;
typedef Eigen::VectorXd VectorXd;

/// 边界条件
Real bnd_ux(Real* p)
{
    Real x = p[0];
    Real y = p[1];
    //return 1 - y * y;
    return 2 * x * y * y * y;
}

Real bnd_uy(Real* p)
{
    Real x = p[0];
    Real y = p[1];
    //return 0;
    return  0.5* pow(x,4) - 0.5 * pow(y, 4);
}

int main(int argc, char* argv[])
{
    int n = 4;
    RectangleDomain * Domain = new RectangleDomain({{-1,-1}, {1,-1}, {1,1}, {-1,1}});
    Mesh<2> * meshU = new Q2Mesh(Domain,{POW2(n), POW2(n)});
    Mesh<2> * meshP = new Q1Mesh(Domain,{POW2(n), POW2(n)});
    Element<2> * Uelement = new Quadrilateral_2_Element();
    Element<2> * Pelement = new Quadrilateral_1_Element();

    int n_u = meshU->n_dofs();
    int n_p = meshP->n_dofs();
    int total_n_dof = 2 * n_u + n_p;
    int n_UeleDof = Uelement -> n_Dofs();
    int n_PeleDof = Pelement -> n_Dofs();
    /// 组装刚度矩阵
    SpMat StiffMat(total_n_dof, total_n_dof);
	SpMat A;
	SpMat Q(n_p,n_p);
    std::vector<Tri> TriList(meshU->n_element() * n_UeleDof * (2*n_UeleDof + 4*n_PeleDof));
	std::vector<Tri> TriListQ(meshP->n_element() * n_PeleDof *n_PeleDof);
    std::vector<Tri>::iterator it = TriList.begin();
	std::vector<Tri>::iterator itq = TriListQ.begin();
    for (int k = 0; k < meshU->n_element(); k++)
    {
	std::vector<int> Idx = meshU->NodeofEle(k);
	std::vector<Dofs<2> > temp(n_UeleDof);
	for (int i = 0; i < n_UeleDof; i++)
	    temp[i] = meshU->DofsofIndex(Idx[i]);
	Uelement->SetDofsList(temp);
	Idx = meshP->NodeofEle(k);
	temp.resize(n_PeleDof);
	for (int i = 0; i < n_PeleDof; i++)
	    temp[i] = meshP->DofsofIndex(Idx[i]);
	Pelement->SetDofsList(temp);

	for (int i = 1; i <= n_UeleDof; i++)
	{
	    for (int j = 1; j <= n_UeleDof; j++)
	    {
		Real a = 0;
		for (int gp = 0; gp < Uelement->n_GaussPnt(); gp++)
		{
		    Real xi = Uelement->GaussionPoint(gp)[0];
		    Real eta = Uelement->GaussionPoint(gp)[1];
		    a += Uelement->det_Jacobi(xi, eta) * Uelement->GaussionWeight(gp) * InnerProuduct(Uelement->gradient(xi, eta, i), Uelement->gradient(xi, eta, j));
		}
		*it = Tri(Uelement->NdIdx(i), Uelement->NdIdx(j), a);
		it++;
		*it = Tri(Uelement->NdIdx(i)+n_u, Uelement->NdIdx(j)+n_u, a);
		it++;
	    }
	    for (int j = 1; j <= n_PeleDof; j++)
	    {
		Real bx = 0;
		Real by = 0;
		for (int gp = 0; gp < Uelement->n_GaussPnt(); gp++)
		{
		    Real xi = Uelement->GaussionPoint(gp)[0];
		    Real eta = Uelement->GaussionPoint(gp)[1];
		    bx += Uelement->det_Jacobi(xi, eta) * Uelement->GaussionWeight(gp) * Pelement->phi(xi, eta, j) * Uelement->phi_x(xi, eta, i);
		    by += Uelement->det_Jacobi(xi, eta) * Uelement->GaussionWeight(gp) * Pelement->phi(xi, eta, j) * Uelement->phi_y(xi, eta, i);
		}
		*it = Tri(Pelement->NdIdx(j)+2*n_u, Uelement->NdIdx(i), -bx);
		it++;
		*it = Tri(Uelement->NdIdx(i), Pelement->NdIdx(j)+2*n_u, -bx);
		it++;
		*it = Tri(Pelement->NdIdx(j)+2*n_u, Uelement->NdIdx(i)+n_u, -by);
		it++;
		*it = Tri(Uelement->NdIdx(i)+n_u, Pelement->NdIdx(j)+2*n_u, -by);
		it++;		
	    }
	}	
		for(int i = 1;i <= n_PeleDof;i++)
		{
			for(int j = 1;j <= n_PeleDof;j++)
			{
				Real q = 0;
				for (int gp = 0; gp < Pelement->n_GaussPnt(); gp++)
				{
		    	Real xi = Pelement->GaussionPoint(gp)[0];
		    	Real eta = Pelement->GaussionPoint(gp)[1];
		    	q += Pelement->det_Jacobi(xi, eta) * Pelement->GaussionWeight(gp) * Pelement->phi(xi,eta,i) * Pelement->phi(xi,eta,j);
				}
				*itq = Tri(Pelement->NdIdx(i),Pelement->NdIdx(j),q);
				itq++;
			}
		}
    }
    StiffMat.setFromTriplets(TriList.begin(), TriList.end());
    StiffMat.makeCompressed();
    Q.setFromTriplets(TriListQ.begin(),TriListQ.end());
	Q.makeCompressed();
	
    /// 处理边界条件
    VectorXd rhs = Eigen::MatrixXd::Zero(total_n_dof, 1);
    for (auto k:meshU->Boundary())
    {
	int bnd_Idxx = k;
	int bnd_Idxy = k + n_u;
	Dofs<2> bnd_point = meshU -> DofsofIndex(k);
	Real bnd_valuex = bnd_ux(*bnd_point);
	Real bnd_valuey = bnd_uy(*bnd_point);
	rhs[bnd_Idxx] = bnd_valuex * StiffMat.coeffRef(bnd_Idxx, bnd_Idxx);
	rhs[bnd_Idxy] = bnd_valuey * StiffMat.coeffRef(bnd_Idxy, bnd_Idxy);
	for (Eigen::SparseMatrix<Real>::InnerIterator it(StiffMat, bnd_Idxx);it;++it)
	{
	    int row = it.row();
	    if (row == bnd_Idxx)
		continue;
	    StiffMat.coeffRef(bnd_Idxx, row) = 0.0;
	    rhs[row] -= StiffMat.coeffRef(row, bnd_Idxx) * bnd_valuex;
	    StiffMat.coeffRef(row, bnd_Idxx) = 0.0;	    
	}
	for (Eigen::SparseMatrix<Real>::InnerIterator it(StiffMat, bnd_Idxy);it;++it)
	{
	    int row = it.row();
	    if (row == bnd_Idxy)
		continue;
	    StiffMat.coeffRef(bnd_Idxy, row) = 0.0;
	    rhs[row] -= StiffMat.coeffRef(row, bnd_Idxy) * bnd_valuey;
	    StiffMat.coeffRef(row, bnd_Idxy) = 0.0;	    
	}
    }
    A = StiffMat.block(0,0,n_u,n_u);
    StokesPreconditioner<double> stokespred(A,Q);
    Eigen::MinRes<SpMat,StokesPreconditioner<double>> solver(stokespred);
    //Eigen::BiCGSTAB<SpMat> solver;
    solver.compute(StiffMat);
    VectorXd x = solver.solve(rhs);
    std::cout << " itertions is :" << solver.iterations() << std::endl;
    std::cout << " error is :" << solver.error() << std::endl;
    VectorXd p = x.tail(n_p);
    Surface<2> S(meshP,Pelement,p);
    S.WriteVTKData3D("Stokes");

    /// 用 matlab 画流线图
    std::cout << "Create matlab file...";
    std::ofstream os;
    os.open("Stokes.m");
    os << "clear;\n";
    double h = meshU->x_h()/2;
    for (int i = 0; i < n_u; i++)
    {
	Dofs<2> temp = meshU->DofsofIndex(i);
	/// 单位化流向向量
	double a = sqrt(x[i] * x[i] + x[i+n_u] * x[i+n_u]) / h;
	os << "quiver( " << temp[0] << "," << temp[1] << "," << x[i]/a << "," << x[i+n_u]/a << ", 'b' "<< ");\n";
	os << "hold on;\n";
    }
    os << "axis equal;\n";
    os.close();
    std::cout<< "     OK!" << std::endl;
    return 0;   
}



